{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[LogicalDevice(name='/device:CPU:0', device_type='CPU'),\n",
       " LogicalDevice(name='/device:GPU:0', device_type='GPU')]"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import gym\n",
    "import matplotlib.pyplot as plt\n",
    "tf.config.list_logical_devices()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.random.set_seed(42)  # extra code – ensures reproducibility on the CPU\n",
    "\n",
    "input_shape = [4]  # == env.observation_space.shape\n",
    "n_outputs = 2  # == env.action_space.n\n",
    "\n",
    "model = tf.keras.Sequential([\n",
    "    tf.keras.layers.Dense(32, activation=\"elu\", input_shape=input_shape),\n",
    "    tf.keras.layers.Dense(32, activation=\"elu\"),\n",
    "    tf.keras.layers.Dense(n_outputs)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def epsilon_greedy_policy(state, epsilon=0):\n",
    "    if np.random.rand() < epsilon:\n",
    "        return np.random.randint(n_outputs)  # random action\n",
    "    else:\n",
    "        Q_values = model.predict(state[np.newaxis], verbose=0)[0]\n",
    "        return Q_values.argmax()  # optimal action according to the DQN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import deque\n",
    "\n",
    "replay_buffer = deque(maxlen=2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_experiences(batch_size):\n",
    "    indices = np.random.randint(len(replay_buffer), size=batch_size)\n",
    "    batch = [replay_buffer[index] for index in indices]\n",
    "    return [\n",
    "        np.array([experience[field_index] for experience in batch])\n",
    "        for field_index in range(6)\n",
    "    ]  # [states, actions, rewards, next_states, dones, truncateds]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def play_one_step(env, state, epsilon):\n",
    "    action = epsilon_greedy_policy(state, epsilon)\n",
    "    next_state, reward, done, truncated, info = env.step(action)\n",
    "    replay_buffer.append((state, action, reward, next_state, done, truncated))\n",
    "    return next_state, reward, done, truncated, info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\"CartPole-v1\", render_mode=\"rgb_array\")\n",
    "env.reset(seed=42)\n",
    "np.random.seed(42)\n",
    "tf.random.set_seed(42)\n",
    "rewards = [] \n",
    "best_score = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "discount_factor = 0.95\n",
    "optimizer = tf.keras.optimizers.Nadam(learning_rate=1e-2)\n",
    "loss_fn = tf.keras.losses.mean_squared_error\n",
    "\n",
    "def training_step(batch_size):\n",
    "    experiences = sample_experiences(batch_size)\n",
    "    states, actions, rewards, next_states, dones, truncateds = experiences\n",
    "    next_Q_values = model.predict(next_states, verbose=0)\n",
    "    max_next_Q_values = next_Q_values.max(axis=1)\n",
    "    runs = 1.0 - (dones | truncateds)  # episode is not done or truncated\n",
    "    target_Q_values = rewards + runs * discount_factor * max_next_Q_values\n",
    "    target_Q_values = target_Q_values.reshape(-1, 1)\n",
    "    mask = tf.one_hot(actions, n_outputs)\n",
    "    with tf.GradientTape() as tape:\n",
    "        all_Q_values = model(states)\n",
    "        Q_values = tf.reduce_sum(all_Q_values * mask, axis=1, keepdims=True)\n",
    "        loss = tf.reduce_mean(loss_fn(target_Q_values, Q_values))\n",
    "\n",
    "    grads = tape.gradient(loss, model.trainable_variables)\n",
    "    optimizer.apply_gradients(zip(grads, model.trainable_variables))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode: 334, Steps: 200, eps: 0.334"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[10], line 5\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[39mfor\u001b[39;00m step \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(\u001b[39m200\u001b[39m):\n\u001b[0;32m      4\u001b[0m     epsilon \u001b[39m=\u001b[39m \u001b[39mmax\u001b[39m(\u001b[39m1\u001b[39m \u001b[39m-\u001b[39m episode \u001b[39m/\u001b[39m \u001b[39m500\u001b[39m, \u001b[39m0.01\u001b[39m)\n\u001b[1;32m----> 5\u001b[0m     obs, reward, done, truncated, info \u001b[39m=\u001b[39m play_one_step(env, obs, epsilon)\n\u001b[0;32m      6\u001b[0m     \u001b[39mif\u001b[39;00m done \u001b[39mor\u001b[39;00m truncated:\n\u001b[0;32m      7\u001b[0m         \u001b[39mbreak\u001b[39;00m\n",
      "Cell \u001b[1;32mIn[6], line 2\u001b[0m, in \u001b[0;36mplay_one_step\u001b[1;34m(env, state, epsilon)\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mplay_one_step\u001b[39m(env, state, epsilon):\n\u001b[1;32m----> 2\u001b[0m     action \u001b[39m=\u001b[39m epsilon_greedy_policy(state, epsilon)\n\u001b[0;32m      3\u001b[0m     next_state, reward, done, truncated, info \u001b[39m=\u001b[39m env\u001b[39m.\u001b[39mstep(action)\n\u001b[0;32m      4\u001b[0m     replay_buffer\u001b[39m.\u001b[39mappend((state, action, reward, next_state, done, truncated))\n",
      "Cell \u001b[1;32mIn[3], line 5\u001b[0m, in \u001b[0;36mepsilon_greedy_policy\u001b[1;34m(state, epsilon)\u001b[0m\n\u001b[0;32m      3\u001b[0m     \u001b[39mreturn\u001b[39;00m np\u001b[39m.\u001b[39mrandom\u001b[39m.\u001b[39mrandint(n_outputs)  \u001b[39m# random action\u001b[39;00m\n\u001b[0;32m      4\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m----> 5\u001b[0m     Q_values \u001b[39m=\u001b[39m model\u001b[39m.\u001b[39;49mpredict(state[np\u001b[39m.\u001b[39;49mnewaxis], verbose\u001b[39m=\u001b[39;49m\u001b[39m0\u001b[39;49m)[\u001b[39m0\u001b[39m]\n\u001b[0;32m      6\u001b[0m     \u001b[39mreturn\u001b[39;00m Q_values\u001b[39m.\u001b[39margmax()\n",
      "File \u001b[1;32me:\\anaconda3\\lib\\site-packages\\keras\\utils\\traceback_utils.py:65\u001b[0m, in \u001b[0;36mfilter_traceback.<locals>.error_handler\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m     63\u001b[0m filtered_tb \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n\u001b[0;32m     64\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m---> 65\u001b[0m     \u001b[39mreturn\u001b[39;00m fn(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m     66\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mException\u001b[39;00m \u001b[39mas\u001b[39;00m e:\n\u001b[0;32m     67\u001b[0m     filtered_tb \u001b[39m=\u001b[39m _process_traceback_frames(e\u001b[39m.\u001b[39m__traceback__)\n",
      "File \u001b[1;32me:\\anaconda3\\lib\\site-packages\\keras\\engine\\training.py:2253\u001b[0m, in \u001b[0;36mModel.predict\u001b[1;34m(self, x, batch_size, verbose, steps, callbacks, max_queue_size, workers, use_multiprocessing)\u001b[0m\n\u001b[0;32m   2251\u001b[0m \u001b[39mfor\u001b[39;00m step \u001b[39min\u001b[39;00m data_handler\u001b[39m.\u001b[39msteps():\n\u001b[0;32m   2252\u001b[0m     callbacks\u001b[39m.\u001b[39mon_predict_batch_begin(step)\n\u001b[1;32m-> 2253\u001b[0m     tmp_batch_outputs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mpredict_function(iterator)\n\u001b[0;32m   2254\u001b[0m     \u001b[39mif\u001b[39;00m data_handler\u001b[39m.\u001b[39mshould_sync:\n\u001b[0;32m   2255\u001b[0m         context\u001b[39m.\u001b[39masync_wait()\n",
      "File \u001b[1;32me:\\anaconda3\\lib\\site-packages\\tensorflow\\python\\util\\traceback_utils.py:150\u001b[0m, in \u001b[0;36mfilter_traceback.<locals>.error_handler\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    148\u001b[0m filtered_tb \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n\u001b[0;32m    149\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m--> 150\u001b[0m   \u001b[39mreturn\u001b[39;00m fn(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m    151\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mException\u001b[39;00m \u001b[39mas\u001b[39;00m e:\n\u001b[0;32m    152\u001b[0m   filtered_tb \u001b[39m=\u001b[39m _process_traceback_frames(e\u001b[39m.\u001b[39m__traceback__)\n",
      "File \u001b[1;32me:\\anaconda3\\lib\\site-packages\\tensorflow\\python\\eager\\def_function.py:915\u001b[0m, in \u001b[0;36mFunction.__call__\u001b[1;34m(self, *args, **kwds)\u001b[0m\n\u001b[0;32m    912\u001b[0m compiler \u001b[39m=\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mxla\u001b[39m\u001b[39m\"\u001b[39m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_jit_compile \u001b[39melse\u001b[39;00m \u001b[39m\"\u001b[39m\u001b[39mnonXla\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[0;32m    914\u001b[0m \u001b[39mwith\u001b[39;00m OptionalXlaContext(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_jit_compile):\n\u001b[1;32m--> 915\u001b[0m   result \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_call(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwds)\n\u001b[0;32m    917\u001b[0m new_tracing_count \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mexperimental_get_tracing_count()\n\u001b[0;32m    918\u001b[0m without_tracing \u001b[39m=\u001b[39m (tracing_count \u001b[39m==\u001b[39m new_tracing_count)\n",
      "File \u001b[1;32me:\\anaconda3\\lib\\site-packages\\tensorflow\\python\\eager\\def_function.py:954\u001b[0m, in \u001b[0;36mFunction._call\u001b[1;34m(self, *args, **kwds)\u001b[0m\n\u001b[0;32m    951\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_lock\u001b[39m.\u001b[39mrelease()\n\u001b[0;32m    952\u001b[0m \u001b[39m# In this case we have not created variables on the first call. So we can\u001b[39;00m\n\u001b[0;32m    953\u001b[0m \u001b[39m# run the first trace but we should fail if variables are created.\u001b[39;00m\n\u001b[1;32m--> 954\u001b[0m results \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_stateful_fn(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwds)\n\u001b[0;32m    955\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_created_variables \u001b[39mand\u001b[39;00m \u001b[39mnot\u001b[39;00m ALLOW_DYNAMIC_VARIABLE_CREATION:\n\u001b[0;32m    956\u001b[0m   \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\u001b[39m\"\u001b[39m\u001b[39mCreating variables on a non-first call to a function\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[0;32m    957\u001b[0m                    \u001b[39m\"\u001b[39m\u001b[39m decorated with tf.function.\u001b[39m\u001b[39m\"\u001b[39m)\n",
      "File \u001b[1;32me:\\anaconda3\\lib\\site-packages\\tensorflow\\python\\eager\\function.py:2496\u001b[0m, in \u001b[0;36mFunction.__call__\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   2493\u001b[0m \u001b[39mwith\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_lock:\n\u001b[0;32m   2494\u001b[0m   (graph_function,\n\u001b[0;32m   2495\u001b[0m    filtered_flat_args) \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_maybe_define_function(args, kwargs)\n\u001b[1;32m-> 2496\u001b[0m \u001b[39mreturn\u001b[39;00m graph_function\u001b[39m.\u001b[39;49m_call_flat(\n\u001b[0;32m   2497\u001b[0m     filtered_flat_args, captured_inputs\u001b[39m=\u001b[39;49mgraph_function\u001b[39m.\u001b[39;49mcaptured_inputs)\n",
      "File \u001b[1;32me:\\anaconda3\\lib\\site-packages\\tensorflow\\python\\eager\\function.py:1862\u001b[0m, in \u001b[0;36mConcreteFunction._call_flat\u001b[1;34m(self, args, captured_inputs, cancellation_manager)\u001b[0m\n\u001b[0;32m   1858\u001b[0m possible_gradient_type \u001b[39m=\u001b[39m gradients_util\u001b[39m.\u001b[39mPossibleTapeGradientTypes(args)\n\u001b[0;32m   1859\u001b[0m \u001b[39mif\u001b[39;00m (possible_gradient_type \u001b[39m==\u001b[39m gradients_util\u001b[39m.\u001b[39mPOSSIBLE_GRADIENT_TYPES_NONE\n\u001b[0;32m   1860\u001b[0m     \u001b[39mand\u001b[39;00m executing_eagerly):\n\u001b[0;32m   1861\u001b[0m   \u001b[39m# No tape is watching; skip to running the function.\u001b[39;00m\n\u001b[1;32m-> 1862\u001b[0m   \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_build_call_outputs(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_inference_function\u001b[39m.\u001b[39;49mcall(\n\u001b[0;32m   1863\u001b[0m       ctx, args, cancellation_manager\u001b[39m=\u001b[39;49mcancellation_manager))\n\u001b[0;32m   1864\u001b[0m forward_backward \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_select_forward_and_backward_functions(\n\u001b[0;32m   1865\u001b[0m     args,\n\u001b[0;32m   1866\u001b[0m     possible_gradient_type,\n\u001b[0;32m   1867\u001b[0m     executing_eagerly)\n\u001b[0;32m   1868\u001b[0m forward_function, args_with_tangents \u001b[39m=\u001b[39m forward_backward\u001b[39m.\u001b[39mforward()\n",
      "File \u001b[1;32me:\\anaconda3\\lib\\site-packages\\tensorflow\\python\\eager\\function.py:499\u001b[0m, in \u001b[0;36m_EagerDefinedFunction.call\u001b[1;34m(self, ctx, args, cancellation_manager)\u001b[0m\n\u001b[0;32m    497\u001b[0m \u001b[39mwith\u001b[39;00m _InterpolateFunctionError(\u001b[39mself\u001b[39m):\n\u001b[0;32m    498\u001b[0m   \u001b[39mif\u001b[39;00m cancellation_manager \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m--> 499\u001b[0m     outputs \u001b[39m=\u001b[39m execute\u001b[39m.\u001b[39;49mexecute(\n\u001b[0;32m    500\u001b[0m         \u001b[39mstr\u001b[39;49m(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49msignature\u001b[39m.\u001b[39;49mname),\n\u001b[0;32m    501\u001b[0m         num_outputs\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_num_outputs,\n\u001b[0;32m    502\u001b[0m         inputs\u001b[39m=\u001b[39;49margs,\n\u001b[0;32m    503\u001b[0m         attrs\u001b[39m=\u001b[39;49mattrs,\n\u001b[0;32m    504\u001b[0m         ctx\u001b[39m=\u001b[39;49mctx)\n\u001b[0;32m    505\u001b[0m   \u001b[39melse\u001b[39;00m:\n\u001b[0;32m    506\u001b[0m     outputs \u001b[39m=\u001b[39m execute\u001b[39m.\u001b[39mexecute_with_cancellation(\n\u001b[0;32m    507\u001b[0m         \u001b[39mstr\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39msignature\u001b[39m.\u001b[39mname),\n\u001b[0;32m    508\u001b[0m         num_outputs\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_num_outputs,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    511\u001b[0m         ctx\u001b[39m=\u001b[39mctx,\n\u001b[0;32m    512\u001b[0m         cancellation_manager\u001b[39m=\u001b[39mcancellation_manager)\n",
      "File \u001b[1;32me:\\anaconda3\\lib\\site-packages\\tensorflow\\python\\eager\\execute.py:54\u001b[0m, in \u001b[0;36mquick_execute\u001b[1;34m(op_name, num_outputs, inputs, attrs, ctx, name)\u001b[0m\n\u001b[0;32m     52\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m     53\u001b[0m   ctx\u001b[39m.\u001b[39mensure_initialized()\n\u001b[1;32m---> 54\u001b[0m   tensors \u001b[39m=\u001b[39m pywrap_tfe\u001b[39m.\u001b[39;49mTFE_Py_Execute(ctx\u001b[39m.\u001b[39;49m_handle, device_name, op_name,\n\u001b[0;32m     55\u001b[0m                                       inputs, attrs, num_outputs)\n\u001b[0;32m     56\u001b[0m \u001b[39mexcept\u001b[39;00m core\u001b[39m.\u001b[39m_NotOkStatusException \u001b[39mas\u001b[39;00m e:\n\u001b[0;32m     57\u001b[0m   \u001b[39mif\u001b[39;00m name \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "for episode in range(600):\n",
    "    obs, info = env.reset()    \n",
    "    for step in range(200):\n",
    "        epsilon = max(1 - episode / 500, 0.01)\n",
    "        obs, reward, done, truncated, info = play_one_step(env, obs, epsilon)\n",
    "        if done or truncated:\n",
    "            break\n",
    "\n",
    "    # extra code – displays debug info, stores data for the next figure, and\n",
    "    #              keeps track of the best model weights so far\n",
    "    print(f\"\\rEpisode: {episode + 1}, Steps: {step + 1}, eps: {epsilon:.3f}\",\n",
    "          end=\"\")\n",
    "    rewards.append(step)\n",
    "    if step >= best_score:\n",
    "        best_weights = model.get_weights()\n",
    "        best_score = step\n",
    "\n",
    "    if episode > 50:\n",
    "        training_step(batch_size)\n",
    "\n",
    "model.set_weights(best_weights)  # extra code – restores the best model weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extra code – this cell generates and saves Figure 18–10\n",
    "plt.figure(figsize=(8, 4))\n",
    "plt.plot(rewards)\n",
    "plt.xlabel(\"Episode\", fontsize=14)\n",
    "plt.ylabel(\"Sum of rewards\", fontsize=14)\n",
    "plt.grid(True)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
